import os
import numpy as np
import torch
import h5py
import json
import hdf5plugin
import shutil
from torch.utils.data import Dataset, DataLoader
from typing import List, Tuple, Dict, Callable, Optional
from rich import print
from rich.progress import track
import cv2

class GenX(Dataset):
    dataset_type = "genx"
    dataset_shape = (10, 20, 360, 640)
    ev_repr_name = "stacked_histogram_dt=50_nbins=10"
    down_sample = 1
    label_columns = ("x", "y", "w", "h", "class_id", "class_confidence")
    classes = {}


    def __init__(self, root: str, mode: str = "train", transform: Optional[Callable] = None) -> None:
        self.root = root
        assert mode in ("train", "val", "test"), "Undefined mode: %s" % (mode,)
        self.mode = mode
        self.transform = transform
        self.data_list = self.get_data_list()


    def get_data_list(self, force_reload: int = False) -> Dict:
        cache_filename = os.path.join(self.root, self.mode + "_cache.json")
        res = dict(
            fnames = [],
            indices = []
        )
        loaded = False
        if os.path.isfile(cache_filename) and not force_reload:
            try:
                with open(cache_filename, "r", encoding = "utf-8") as f:
                    res = json.load(f)
                loaded = True
            except:
                loaded = False
        if not loaded:
            file_list = os.listdir(os.path.join(self.root, self.mode))
            res ["fnames"] = file_list
            for fidx, fname in enumerate(file_list):
                samples = np.load(os.path.join(self.ev_repr_folder(fname), "objframe_idx_2_repr_idx.npy")).tolist()
                for tidx, ts in enumerate(samples):
                    if ts < self.dataset_shape[0]:
                        continue
                    res["indices"].append([fidx, tidx])
            with open(cache_filename, "w", encoding = "utf-8") as f:
                json.dump(res, f)
        return res


    def ev_repr_folder(self, idx: str) -> str:
        return os.path.join(self.root, self.mode, idx, "event_representations_v2", self.ev_repr_name)


    @property
    def ev_repr_filename(self) -> str:
        return "event_representations%s.h5" % (("_ds%d_nearest" % self.down_sample) if self.down_sample > 1 else "")


    def label_folder(self, idx: str) -> str:
        return os.path.join(self.root, self.mode, idx, "labels_v2")


    def get_data(self, idx: Tuple[int, int]) -> Tuple[torch.Tensor, torch.Tensor]:
        fidx, tidx = idx
        data_dir = self.data_list["fnames"][fidx]
        events_filename = os.path.join(self.ev_repr_folder(data_dir), self.ev_repr_filename)
        events_idx_filename = os.path.join(self.ev_repr_folder(data_dir), "objframe_idx_2_repr_idx.npy")

        objframe_idx_to_data_idx = np.load(events_idx_filename)
        end_t = objframe_idx_to_data_idx[tidx] + 1
        start_t = end_t - self.dataset_shape[0]
        with h5py.File(events_filename, 'r') as h5f:
            data = torch.from_numpy(h5f['data'][start_t:end_t])
        return data

    def get_labels(self, idx: Tuple[int, int]) -> Tuple[torch.Tensor, torch.Tensor]:
        fidx, tidx = idx
        data_dir = self.data_list["fnames"][fidx]
        labels_filename = os.path.join(self.label_folder(data_dir), "labels.npz")

        label_data = np.load(labels_filename)
        objframe_idx_to_label_idx = label_data["objframe_idx_2_label_idx"]
        np_labels = label_data['labels']

        start_idx = objframe_idx_to_label_idx[tidx]
        end_idx = objframe_idx_to_label_idx[tidx + 1] if (tidx < len(objframe_idx_to_label_idx) - 1) else None
        labels = torch.tensor(np.array([np_labels[k][start_idx:end_idx].astype(np.float64) for k in self.label_columns]))
        labels = labels.permute(1, 0).to(torch.double)
        labels[:, :4:2] /= self.dataset_shape[-1]
        labels[:, 1:4:2] /= self.dataset_shape[-2]
        labels[:, :2] += labels[:, 2:4] / 2
        return labels


    def __len__(self):
        return len(self.data_list["indices"])


    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        idx_slice = tuple(self.data_list["indices"][idx])
        data = self.get_data(idx_slice)
        labels = self.get_labels(idx_slice)
        if self.transform is not None:
            data = self.transform(data)
        return data, labels


class Gen1(GenX):
    dataset_type = "gen1"
    dataset_shape = (10, 20, 240, 304)
    down_sample = 1
    classes = {
        0: "car",
        1: "pedestrian"
    }


class Gen4(GenX):
    dataset_type = "gen4"
    dataset_shape = (4, 20, 360, 640)
    down_sample = 2
    classes = {
        0: "pedestrian",
        1: "two-wheeler",
        2: "car"
    }


def genx_collate(raw_data: List[Tuple[torch.Tensor, torch.Tensor]]) -> Dict:
    res_idx = []
    res_data = []
    res_bbox = []
    res_label = []
    for idx, (x, y) in enumerate(raw_data):
        res_idx.extend([idx] * len(y))
        res_data.append(x)
        res_bbox.append(y[:, :4])
        res_label.append(y[:, 4:5])

    return {
        'im_file': None,
        'ori_shape': [(x.shape[-2], x.shape[-1])] * len(raw_data),
        'resized_shape': [(x.shape[-2], x.shape[-1])] * len(raw_data),
        'img': torch.stack(res_data).permute(1, 0, 2, 3, 4),
        'cls': torch.cat(res_label, dim = 0).to(torch.float),
        'bboxes': torch.cat(res_bbox, dim = 0).to(torch.double),
        'batch_idx': torch.tensor(res_idx, dtype = torch.float32),
        'ratio_pad': [((1, 1), (0, 0))] * len(raw_data)
    }